package com.scheduler.utils.trace;

import com.scheduler.utils.thread.ThreadUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicInteger;


public class Slf4jMdcUtil {
    private static final Logger LOGGER = LoggerFactory.getLogger(Slf4jMdcUtil.class);

    /**
     * 存储进入的次数，进入次数减为0的时候才clear
     */
    private static final ThreadLocal<AtomicInteger> ENTRANCE_TIMES = ThreadLocal.withInitial(AtomicInteger::new);

    /**
     * 全新追加一个，否则可能被上一个线程污染
     */
    public static void appendNewTraceIdMDC() {
        appendMDC(null, null, true);
    }

    /**
     * 这个方法应该仅在主线程往子线程同步traceId的情况下使用
     */
    public static void appendMDCWithTraceId(String traceId) {
        appendMDC(traceId, null, false);
    }

    /**
     * @deprecated use {@link #appendNewTraceIdMDC()} and {@link #appendMDCWithTraceId(String)}
     */
    @Deprecated
    public static void appendMDC() {
        appendMDC(null, null, false);
    }

    public static void appendMDC(String traceId, String bizChainId) {
        appendMDC(traceId, bizChainId, false);
    }

    public static void appendMDC(String traceId, String bizChainId, boolean isNew) {
        try {
            if (isNew) {
                // 先清理
                clearMDCAndComplete();
            }
            if (ENTRANCE_TIMES.get().getAndIncrement() > 0) {
                // 重入的直接跳过（先 get 再 +1 ）
                return;
            }
            // 如果当前线程中已经有traceId，不用再执行一遍
            if (isNew) {
                traceId = TraceUtil.createTraceId();
            } else {
                traceId = TraceUtil.getOrCreateTraceId(traceId);
            }

            // 兼容QA的日志%X{traceId}配置
            MDC.put("traceId", traceId);
            MDC.put("TRACE_ID", traceId);
            if (bizChainId != null) {
                // 只在入口写一次，后续不覆盖。BIZ_CHAIN_ID 的定义是，识别一个traceId中的某一个链路，因为一个traceId会很长，而链路只关注当前链路本身。非必须，仅traceId非常长的情况下才有效果，按需增加。
                MDC.put("BIZ_CHAIN_ID", bizChainId);
            }
        } catch (Throwable t) {
            LOGGER.error("添加MDC异常", t);
        }
    }

    public static void appendCustomMDC(String key, String value) {
        try {
            MDC.put(key, value);
        } catch (Throwable t) {
            LOGGER.error("appendCustomMDC error {}:{}", key, value, t);
        }
    }

    public static void clearMDC() {
        try {
            int i = ENTRANCE_TIMES.get().decrementAndGet();
            if (i > 0) {
                return;
            }
            clearMDCAndComplete();
        } catch (Throwable t) {
            LOGGER.error("清理MDC异常", t);
        }
    }

    /**
     * 直接、完全清理，避免出现重入的清理不完全导致线程污染
     */
    public static void clearMDCAndComplete() {
        try {
            // threadLocal也清除掉
            ENTRANCE_TIMES.remove();

            // 兼容QA的日志%X{traceId}配置
            MDC.remove("traceId");
            MDC.remove("TRACE_ID");
            MDC.remove("TENANT_ID");
            MDC.remove("BIZ_CODE");
            MDC.remove("BIZ_CHAIN_ID");
            MDC.remove("tenantId");
            MDC.remove("virtualRegion");
            MDC.remove("uid");

            // 完成trace
            TraceUtil.endTraceId();
        } catch (Throwable t) {
            LOGGER.error("清理MDC异常2", t);
        }
    }

    /**
     * 查看当前线程的MDC重入次数
     */
    public static int currentTraceRetryTimes() {
        return ENTRANCE_TIMES.get().get();
    }

    /**
     * 只有调用了 appendMDC，在这里才能读取到，否则读取不到
     */
    public static String traceId() {
        return MDC.get("traceId");
    }



    /**
     * 包装正常执行的runnable，主要实现租户ID、traceId从主线程写入到子线程
     *
     * @param runnable 正常执行的runnable
     * @return 包装后的runnable
     */
    public static Runnable wrapperRunnableForSubThread(Runnable runnable) {
        return ThreadUtil.wrapperRunnableForSubThread(runnable);
    }

    /**
     * 包装正常执行的callable，主要实现租户ID、traceId从主线程写入到子线程
     *
     * @param callable 正常执行的callable
     * @return 包装后的callable
     */
    public static <V> Callable<V> wrapperCallableForSubThread(Callable<V> callable) {
        return ThreadUtil.wrapperCallableForSubThread(callable);
    }
}

